import torch
from modelscope.utils.constant import Tasks
from modelscope.pipelines import pipeline
from modelscope.preprocessors.image import load_image
from fastapi import WebSocket

class multimModaImageTextSearch:
    # AI模型的本地地址
    modelName: str = 'damo/multi-modal_clip-vit-large-patch14_336_zh'
    # 使用 cpu 还是 GPU 计算，GPU 比 CPU快，但是 GPU 只支持 英伟达 显卡，至少4G显存以上
    deviceType: str = 'cpu' # cpu, gpu
    # 干扰词，干扰词越少见，搜索结果的准确度越高
     # 但不是越少见越多就越好，需要用数据测试，找到合适的
    noiseWord: list[str] = ['小火龙']
    # 阈值，数值越大，搜索结果的准确度越高，
    # 但不是越大越好，需要用数据测试，找到合适的
    thresholdValue: int = 85
    #  外部传入 需要检索的关键词数组
    keyword: str = ""
    # 外部传入 的视频帧 截图数据，[ 'http地址 或者 本地地址', '' ]
    # 地址可以用 video标签 poster 预览图，或者使用 ffmpeg，docarray，这里推荐 
    # 使用 多种模态数据结构工具包 docarray，具体 docarray 封装看 utils/screenshot.py
    imgArr: list[str] = []
    # 符合 条件 数据的计数
    count: int = 0
    # 符合 条件 数据的list下标
    haveIndex = []

    def __init__(self, key, imgList):
        self.keyword = key
        self.imgArr = imgList

    def start(self):
        self.count = 0
        self.haveIndex = []

        for index, item in enumerate(self.imgArr):
            print()
            print("关键词: ", self.keyword)
            print(f"下标：{index}，图片地址为：{item}")
            res = self.imageTextSearch(item)
            print(f"下标：{index}，图片&文字是否相似：{res}")
            if res:
                self.count += 1
                self.haveIndex.append(index)

        if self.count != 0 and len(self.haveIndex) != 0:
            return {'count': self.count, 'index': self.haveIndex}
        else:
            return {'count': 0, 'index': []}
        
    def imageTextSearch(self,url):
        pipeline_caption = pipeline(
            task=Tasks.multi_modal_embedding, 
            model=self.modelName, 
            model_revision='v1.0.1', 
            device=self.deviceType
        )
        # 支持网络地址 & 本地地址
        input_img = load_image(url) 
        # input_texts = [self.noiseWord, self.keyword]
        input_texts = [self.keyword]+[item for item in self.noiseWord]

        img_embedding = pipeline_caption.forward({'img': input_img})['img_embedding'] 
        text_embedding = pipeline_caption.forward({'text': input_texts})['text_embedding']

        with torch.no_grad():
            logits_per_image = (img_embedding / pipeline_caption.model.temperature) @ text_embedding.t()
            probs = logits_per_image.softmax(dim=-1).cpu().numpy()

        similarityNum = probs[0][0] * 100
        print(f"图片&文字相似度数值：{similarityNum}")
        result = True if similarityNum >= self.thresholdValue else False
        return result
              